{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "import xml.etree.ElementTree as ET\n",
    "import os"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def convert(size, box):\n",
    "    dw = 1. / size[0]\n",
    "    dh = 1. / size[1]\n",
    "    x = (box[0] + box[1]) / 2.0 - 1\n",
    "    y = (box[2] + box[3]) / 2.0 - 1\n",
    "    w = box[1] - box[0]\n",
    "    h = box[3] - box[2]\n",
    "    x = x * dw\n",
    "    w = w * dw\n",
    "    y = y * dh\n",
    "    h = h * dh\n",
    "    return (x, y, w, h)\n",
    "\n",
    "def convert_annotation(xml_file, txt_file, classes):\n",
    "    tree = ET.parse(xml_file)\n",
    "    root = tree.getroot()\n",
    "    size = root.find('size')\n",
    "    w = int(size.find('width').text)\n",
    "    h = int(size.find('height').text)\n",
    "\n",
    "    with open(txt_file, 'w') as out_file:\n",
    "        for obj in root.iter('object'):\n",
    "            difficult = obj.find('difficult').text\n",
    "            cls = obj.find('name').text\n",
    "            if cls not in classes or int(difficult) == 1:\n",
    "                continue\n",
    "            cls_id = classes.index(cls)\n",
    "            xmlbox = obj.find('bndbox')\n",
    "            b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text),\n",
    "                 float(xmlbox.find('ymin').text), float(xmlbox.find('ymax').text))\n",
    "            bb = convert((w, h), b)\n",
    "            out_file.write(f\"{cls_id} {' '.join([str(a) for a in bb])}\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "classes = [\"site\",\"wall\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# xml_file = './xml_label/1.xml'\n",
    "# txt_file = './txt_label/1.txt'\n",
    "# convert_annotation(xml_file, txt_file, classes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "ename": "IsADirectoryError",
     "evalue": "[Errno 21] Is a directory: './xml_label'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mIsADirectoryError\u001b[0m                         Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[10], line 10\u001b[0m\n\u001b[1;32m      8\u001b[0m         convert_annotation(xml_file, txt_file, classes)\n\u001b[1;32m      9\u001b[0m txt_file \u001b[38;5;241m=\u001b[39m \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mpath/to/your/output.txt\u001b[39m\u001b[38;5;124m'\u001b[39m\n\u001b[0;32m---> 10\u001b[0m \u001b[43mconvert_annotation\u001b[49m\u001b[43m(\u001b[49m\u001b[43mxml_folder\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtxt_folder\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mclasses\u001b[49m\u001b[43m)\u001b[49m\n",
      "Cell \u001b[0;32mIn[7], line 15\u001b[0m, in \u001b[0;36mconvert_annotation\u001b[0;34m(xml_file, txt_file, classes)\u001b[0m\n\u001b[1;32m     14\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mconvert_annotation\u001b[39m(xml_file, txt_file, classes):\n\u001b[0;32m---> 15\u001b[0m     tree \u001b[38;5;241m=\u001b[39m \u001b[43mET\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mparse\u001b[49m\u001b[43m(\u001b[49m\u001b[43mxml_file\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     16\u001b[0m     root \u001b[38;5;241m=\u001b[39m tree\u001b[38;5;241m.\u001b[39mgetroot()\n\u001b[1;32m     17\u001b[0m     size \u001b[38;5;241m=\u001b[39m root\u001b[38;5;241m.\u001b[39mfind(\u001b[38;5;124m'\u001b[39m\u001b[38;5;124msize\u001b[39m\u001b[38;5;124m'\u001b[39m)\n",
      "File \u001b[0;32m/usr/lib/python3.11/xml/etree/ElementTree.py:1218\u001b[0m, in \u001b[0;36mparse\u001b[0;34m(source, parser)\u001b[0m\n\u001b[1;32m   1209\u001b[0m \u001b[38;5;250m\u001b[39m\u001b[38;5;124;03m\"\"\"Parse XML document into element tree.\u001b[39;00m\n\u001b[1;32m   1210\u001b[0m \n\u001b[1;32m   1211\u001b[0m \u001b[38;5;124;03m*source* is a filename or file object containing XML data,\u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m   1215\u001b[0m \n\u001b[1;32m   1216\u001b[0m \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m   1217\u001b[0m tree \u001b[38;5;241m=\u001b[39m ElementTree()\n\u001b[0;32m-> 1218\u001b[0m \u001b[43mtree\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mparse\u001b[49m\u001b[43m(\u001b[49m\u001b[43msource\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mparser\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1219\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m tree\n",
      "File \u001b[0;32m/usr/lib/python3.11/xml/etree/ElementTree.py:569\u001b[0m, in \u001b[0;36mElementTree.parse\u001b[0;34m(self, source, parser)\u001b[0m\n\u001b[1;32m    567\u001b[0m close_source \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mFalse\u001b[39;00m\n\u001b[1;32m    568\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28mhasattr\u001b[39m(source, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mread\u001b[39m\u001b[38;5;124m\"\u001b[39m):\n\u001b[0;32m--> 569\u001b[0m     source \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mopen\u001b[39m(source, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mrb\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m    570\u001b[0m     close_source \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mTrue\u001b[39;00m\n\u001b[1;32m    571\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n",
      "\u001b[0;31mIsADirectoryError\u001b[0m: [Errno 21] Is a directory: './xml_label'"
     ]
    }
   ],
   "source": [
    "xml_folder = './xml_label'\n",
    "txt_folder = './txt_label'\n",
    "\n",
    "for file in os.listdir(xml_folder):\n",
    "    if file.endswith('.xml'):\n",
    "        xml_file = os.path.join(xml_folder, file)\n",
    "        txt_file = os.path.join(txt_folder, file.replace('.xml', '.txt'))\n",
    "        convert_annotation(xml_file, txt_file, classes)\n",
    "txt_file = './txt_label'\n",
    "convert_annotation(xml_folder, txt_folder, classes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.0rc1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
